"""内网穿透客户端"""
import asyncio
import json
import logging
from base64 import b64decode, b64encode
import aiohttp
from homeassistant.helpers.aiohttp_client import async_get_clientsession
from homeassistant.core import HomeAssistant

_LOGGER = logging.getLogger(__name__)

class TunnelClient:
    """内网穿透客户端"""
    
    def __init__(self, hass: HomeAssistant, server_url: str, username: str, password: str):
        self.hass = hass
        self.server_url = server_url
        self.username = username
        self.password = password
        self.ws_client = None
        self.running = False
        self.reconnect_task = None
        self.local_ws_connections = {}
    
    async def connect(self):
        """连接到隧道服务器"""
        if self.running:
            return
            
        self.running = True
        self.reconnect_task = self.hass.async_create_task(self._reconnect_loop())

    async def _reconnect_loop(self):
        """重连循环"""
        while self.running:
            try:
                await self._connect_websocket()
                if self.ws_client and not self.ws_client.closed:
                    _LOGGER.info("隧道客户端连接成功")
                    await self._handle_messages()
            except asyncio.CancelledError:
                break
            except Exception as e:
                if self.running:
                    _LOGGER.debug(f"隧道连接失败，30秒后重试: {e}")
                    
            if self.running:
                try:
                    await asyncio.sleep(30)
                except asyncio.CancelledError:
                    break

    async def _connect_websocket(self):
        """连接WebSocket"""
        session = async_get_clientsession(self.hass)
        url = f"{self.server_url}?username={self.username}&password={self.password}"
        self.ws_client = await session.ws_connect(url)

    async def _handle_messages(self):
        """处理消息"""
        try:
            async for msg in self.ws_client:
                if msg.type == aiohttp.WSMsgType.TEXT:
                    try:
                        data = json.loads(msg.data)
                        await self._process_message(data)
                    except Exception as e:
                        _LOGGER.error(f"处理消息错误: {e}")
                elif msg.type in (aiohttp.WSMsgType.CLOSED, aiohttp.WSMsgType.ERROR):
                    break
        except Exception as e:
            _LOGGER.debug(f"消息处理循环错误: {e}")

    async def _process_message(self, data):
        """处理具体消息"""
        msg_type = data.get("type")
        
        if msg_type == "connected":
            _LOGGER.info(f"隧道已建立，子域名: {data.get('subdomain')}")
        elif msg_type == "http_request":
            await self._handle_http_request(data)
        elif msg_type == "ws_connect":
            await self._handle_ws_connect(data)
        elif msg_type == "ws_data":
            await self._handle_ws_data(data)
        elif msg_type == "ws_close":
            await self._handle_ws_close(data)

    async def _handle_http_request(self, data):
        """处理HTTP请求"""
        request_id = data["requestId"]
        method = data["method"]
        url = data["url"]
        headers = data["headers"]
        body = b64decode(data["body"]) if data["body"] else None
        
        try:
            session = async_get_clientsession(self.hass)
            local_url = f"http://127.0.0.1:1880{url}"
            local_headers = {k: v for k, v in headers.items() if k.lower() != 'host'}
            
            # POST请求延长超时时间
            timeout_seconds = 60 if method == 'POST' else 30
            async with session.request(
                method, local_url, headers=local_headers, data=body,
                timeout=aiohttp.ClientTimeout(total=timeout_seconds)
            ) as response:
                resp_body = await response.read()
                
                # 分包处理大响应
                max_chunk_size = 32 * 1024  # 64KB
                if len(resp_body) > max_chunk_size:
                    # 发送分包开始
                    await self._send_response({
                        "type": "http_chunk_start",
                        "requestId": request_id,
                        "statusCode": response.status,
                        "headers": dict(response.headers)
                    })
                    
                    # 分包发送
                    for i in range(0, len(resp_body), max_chunk_size):
                        chunk = resp_body[i:i + max_chunk_size]
                        await self._send_response({
                            "type": "http_chunk",
                            "requestId": request_id,
                            "data": b64encode(chunk).decode('utf-8')
                        })
                        await asyncio.sleep(0.001)  # 避免阻塞
                    
                    # 发送结束
                    await self._send_response({
                        "type": "http_chunk_end",
                        "requestId": request_id
                    })
                else:
                    # 小响应直接发送
                    await self._send_response({
                        "type": "http_response",
                        "requestId": request_id,
                        "statusCode": response.status,
                        "headers": dict(response.headers),
                        "body": b64encode(resp_body).decode('utf-8') if resp_body else ""
                    })
                
        except Exception as e:
            await self._send_response({
                "type": "http_response",
                "requestId": request_id,
                "statusCode": 500,
                "headers": {},
                "body": b64encode(str(e).encode('utf-8')).decode('utf-8')
            })

    async def _handle_ws_connect(self, data):
        """处理WebSocket连接"""
        request_id = data["requestId"]
        url = data["url"]
        headers = data["headers"]
        
        try:
            session = async_get_clientsession(self.hass)
            local_url = f"ws://127.0.0.1:1880{url}"
            local_headers = {k: v for k, v in headers.items() if k.lower() != 'host'}
            
            local_ws = await session.ws_connect(local_url, headers=local_headers)
            self.local_ws_connections[request_id] = local_ws
            
            # 启动转发任务
            self.hass.async_create_task(
                self._forward_ws_messages(local_ws, request_id)
            )
            
        except Exception as e:
            _LOGGER.error(f"WebSocket连接失败: {e}")

    async def _forward_ws_messages(self, local_ws, request_id):
        """转发本地WebSocket消息"""
        try:
            async for msg in local_ws:
                if request_id not in self.local_ws_connections:
                    break
                    
                if msg.type == aiohttp.WSMsgType.TEXT:
                    await self._send_response({
                        "type": "ws_data",
                        "requestId": request_id,
                        "data": msg.data
                    })
                elif msg.type == aiohttp.WSMsgType.BINARY:
                    await self._send_response({
                        "type": "ws_data",
                        "requestId": request_id,
                        "data": msg.data.decode('utf-8', errors='ignore')
                    })
                elif msg.type in (aiohttp.WSMsgType.CLOSED, aiohttp.WSMsgType.ERROR):
                    break
        except Exception:
            pass
        finally:
            if request_id in self.local_ws_connections:
                del self.local_ws_connections[request_id]

    async def _handle_ws_data(self, data):
        """处理WebSocket数据"""
        request_id = data["requestId"]
        message_data = data["data"]
        
        if request_id in self.local_ws_connections:
            local_ws = self.local_ws_connections[request_id]
            try:
                await local_ws.send_str(message_data)
            except Exception:
                pass

    async def _handle_ws_close(self, data):
        """处理WebSocket关闭"""
        request_id = data["requestId"]
        if request_id in self.local_ws_connections:
            local_ws = self.local_ws_connections[request_id]
            if not local_ws.closed:
                await local_ws.close()
            del self.local_ws_connections[request_id]

    async def _send_response(self, message):
        """发送响应到服务器"""
        if self.ws_client and not self.ws_client.closed:
            try:
                await self.ws_client.send_json(message)
            except Exception:
                pass

    async def disconnect(self):
        """断开连接"""
        self.running = False
        
        if self.reconnect_task:
            self.reconnect_task.cancel()
            try:
                await self.reconnect_task
            except asyncio.CancelledError:
                pass
            
        if self.ws_client and not self.ws_client.closed:
            await self.ws_client.close()
            
        for ws in self.local_ws_connections.values():
            if not ws.closed:
                await ws.close()
        self.local_ws_connections.clear()